package com.zsn.ai.config;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.conditions.query.QueryChainWrapper;
import com.zsn.ai.entity.po.ChatMessageHistory;
import com.zsn.ai.entity.po.ChatidHistory;
import com.zsn.ai.mapper.ChatMessageHistoryMapper;
import com.zsn.ai.service.IChatMessageHistoryService;
import com.zsn.ai.service.IChatidHistoryService;
import java.util.ArrayList;
import java.util.List;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

/**
 * @program: ai-demo
 * @ClassName MysqlChatMemoryRepository
 * @description:
 * @author: zsn
 * @create: 2025-10-02 00:37
 * @Version 1.0
 **/
//@Component
public class MysqlChatMemoryRepository implements ChatMemoryRepository {

    @Autowired
    private IChatMessageHistoryService chatMessageHistoryService;
    @Autowired
    private IChatidHistoryService chatidHistoryService;
    @Override
    public List<String> findConversationIds() {
        List<ChatidHistory> list = chatidHistoryService.list();
        return list.stream().map(ChatidHistory::getChatId).toList();
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        LambdaQueryWrapper<ChatMessageHistory> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(ChatMessageHistory::getChatId, conversationId);
        List<ChatMessageHistory> list = chatMessageHistoryService.list(wrapper);
        List<Message> messages = new ArrayList<>();
        for (ChatMessageHistory chatMessage : list) {
            switch (chatMessage.getRole()) {
                case "user":
                    messages.add(new UserMessage(chatMessage.getContent()));
                    break;
                case "assistant":
                    messages.add(new AssistantMessage(chatMessage.getContent()));
                    break;
                case "system":
                    messages.add(new SystemMessage(chatMessage.getContent()));
                    break;
                default:
                    throw new IllegalArgumentException("Unknown role: " + chatMessage.getRole());
            }

        }
        return messages;
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        List<ChatMessageHistory> chatMessages = new ArrayList<>();
        for (Message message : messages) {
            String role = "";
            switch (message.getMessageType()) {
                case USER:
                    role = "user";
                    break;
                case ASSISTANT:
                    role = "assistant";
                    break;
                case SYSTEM:
                    role = "system";
                    break;
                default:
                    role = "unknown";
                    break;
            }
            ChatMessageHistory chatMessage = new ChatMessageHistory();
            chatMessage.setChatId(conversationId);
            chatMessage.setRole(role);
            chatMessage.setContent(message.getText());
//            chatMessageHistoryService.save(chatMessage);
            chatMessages.add(chatMessage);
        }
        // 插入到数据库
        chatMessageHistoryService.saveBatch(chatMessages);
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        LambdaQueryWrapper<ChatMessageHistory> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(ChatMessageHistory::getChatId, conversationId);
        chatMessageHistoryService.remove(wrapper);
    }
}